//
// Event.cxx
//
// Created by TB on 11/22/11.
// Copyright 2011 BNL. All rights reserved.
//

#include <algorithm>
#include <cmath>
#include <cstdlib>
#include <functional>
#include <iostream>

#include <TCollection.h> // For TIter
#include <TLorentzVector.h>
#include <TMath.h>
#include <TMCParticle.h>
#include <TObjArray.h>
#include <TPythia6.h>
#include <TRandom3.h>
#include <TRotation.h>
#include <TVector3.h>

#include "TParticle.h"
#include "TPDGCode.h"
#include "TParticlePDG.h"

#include <boost/foreach.hpp>

#include "tppmc/Distributions.h"
#include "tppmc/Enumerations.h"
#include "tppmc/Event.h"
#include "tppmc/FragmentationFunction.h"
#include "tppmc/Particle.h"
#include "tppmc/PartonContent.h"
#include "tppmc/RandomSpin.h"

namespace {
  bool isParton(int pdg) {
    return abs(pdg) < 7 or 21 == pdg;
  }

  // Helper functor for sorting Partons by group
  struct SortPartonsByGroup {
    bool operator()(const tppmc::Parton& a, const tppmc::Parton& b) {
      return a.Group() < b.Group();
    }
  };
}

namespace tppmc {

  Event::Event(double rotation, double etamin, double etamax, double minpt, long tparticleout)
    : mRotation(rotation)
    , mEtaMin(etamin)
    , mEtaMax(etamax)
    , mMinPt(minpt)
      // Need full namespace specification in TClonesArray constructor.
    , mPartons("tppmc::Parton")
    , mHadrons("tppmc::Hadron")
    , mParticle("TParticle") 
    , mtparticleout(tparticleout)
  {
  }

  unsigned Event::Update(Distributions& pdfs) {
      
    // Access the current PYTHIA state via the singleton TPythia6.
    TPythia6* pythia = TPythia6::Instance();
      
    mScatter = std::vector<Parton>(2, Parton());
      
    mProcessId = pythia->GetMSTI(1);
    mQSquared = pythia->GetPARI(22);
    mPartonPt = pythia->GetPARI(17);

    ComputeCentreOfMassEnergy();
    FillParticles(pdfs);
      
    int ninitial = mParticle.GetEntries();
    int nhadrons = Fragment(*(pdfs.Fragmentation()));
    ninitial -= mParticle.GetEntries();
    ninitial = abs(ninitial);
    if (mtparticleout==1){
      mPartons.Clear();
    }
    return std::max(nhadrons,ninitial);

  } 

  void Event::ComputeCentreOfMassEnergy() {
    // Access the current PYTHIA state via the singleton TPythia6.
    TObjArray* particles = TPythia6::Instance()->ImportParticles();
    // Get the first 2 particles, which are the colliding beams
    TMCParticle* beam1 = dynamic_cast<TMCParticle*>(particles->At(0));
    TMCParticle* beam2 = dynamic_cast<TMCParticle*>(particles->At(1));
    // Use 4-vector class to compute centre-of-mass energy.
    TLorentzVector p1(beam1->GetPx(), beam1->GetPy(), beam1->GetPz(),
		      beam1->GetEnergy());
    TLorentzVector p2(beam2->GetPx(), beam2->GetPy(), beam2->GetPz(),
		      beam2->GetEnergy());
    mCentreOfMassEnergy = (p1 + p2).Mag();
  }

  void Event::FillParticles(Distributions& pdfs) {
    mPartons.Clear();
    mHadrons.Clear();
    mParticle.Clear();
    // Access the current PYTHIA state via the singleton TPythia6.
    TObjArray* particles = TPythia6::Instance()->ImportParticles();
    // Loop over particles and extract partons and J/Psi.
    TIter next(particles);
    TMCParticle* particle(NULL);
    unsigned counter(0);
    // List to accumulate all partons.
    // We store them so we can sort them by group, so group 1 partons
    // precede group 2 partons, before adding them to the event.
    std::list<Parton> partons;
    int      pspin[2]={0,0};
    while((particle = dynamic_cast<TMCParticle*>(next()))) {
      Int_t    type       = particle->GetKF();
      Int_t    status     = particle->GetKS();
      Int_t    parent     = particle->GetParent() - 1;
      Double_t px         = particle->GetPx();
      Double_t py         = particle->GetPy();
      Double_t pz         = particle->GetPz();
      // Event plane rotation.
      // mRotation was set externally before calling Update().
      Double_t pyth_pt  = sqrt(px * px + py * py); 
      Double_t pyth_phi = atan2(py, px);
      px = pyth_pt * cos(pyth_phi + mRotation);
      py = pyth_pt * sin(pyth_phi + mRotation);
      if ( counter == 0 ) {//first proton            
	int spin = gRandom->Integer(2);
	if(spin == 0) {
	  pspin[0] = -1; 
	} // if
	else {
	  pspin[0] = 1;
	} // if
      }
      if ( counter == 1 ) {//first proton            
	int spin = gRandom->Integer(2);
	if(spin == 0) {
	  pspin[1] = -1; 
	} // if
	else {
	  pspin[1] = 1;
	} // if
      }
      // Compute Sivers weights from the partons involved
      // in the initial hard interaction
      // First parton in hard interaction has index 4 in the particle list
      // Second parton in hard interaction has index 5 in the particle list
      if(4 == counter) {
	// this can go in a function
	mScatter.at(0) = Parton(px, py, pz, type, 1);
	mScatter.at(0).x = TPythia6::Instance()->GetPARI(33);
	mScatter.at(0).mSiversWeight = ComputeSiversWeight(mScatter.at(0), pdfs);
	// Note the minus sign
	int lhapdf = mScatter.at(0).Lhapdf();
	PartonContent values;
	pdfs.Unpolarised()->Evaluate(mScatter.at(0), mQSquared, values);
	mScatter.at(0).mF1 = values[lhapdf];
	pdfs.Sivers()->Evaluate(mScatter.at(0), mQSquared, values);
	mScatter.at(0).mSivers = values[lhapdf];
	pdfs.Transversity()->Evaluate(mScatter.at(0), mQSquared, values);
	mScatter.at(0).mH1 = values[lhapdf];
	mScatter.at(0).pspin = pspin[0];
	mScatter.at(0).mSiversSpinFactor = ComputeSiversSpinFactor(mScatter.at(0));
	mScatter.at(0).spin =
	  mScatter.at(0).pspin * ComputeTransverseSpin(mScatter.at(0), pdfs);
      } // if
      else if(5 == counter) {
	mScatter.at(1) = Parton(px, py, pz, type, 2);
	mScatter.at(1).x = TPythia6::Instance()->GetPARI(34);
	mScatter.at(1).mSiversWeight = ComputeSiversWeight(mScatter.at(1), pdfs);
	int lhapdf = mScatter.at(1).Lhapdf();
	PartonContent values;
	pdfs.Unpolarised()->Evaluate(mScatter.at(1), mQSquared, values);
	mScatter.at(1).mF1 = values[lhapdf];
	pdfs.Sivers()->Evaluate(mScatter.at(1), mQSquared, values);
	mScatter.at(1).mSivers = values[lhapdf];
	pdfs.Transversity()->Evaluate(mScatter.at(1), mQSquared, values);
	mScatter.at(1).mH1 = values[lhapdf];
	mScatter.at(1).pspin = pspin[1];
	mScatter.at(1).mSiversSpinFactor = -ComputeSiversSpinFactor(mScatter.at(1));
	mScatter.at(1).spin =
	  mScatter.at(1).pspin * ComputeTransverseSpin(mScatter.at(1), pdfs);
      } // else if
         
      if(isParton(type)) {
	if(6 == parent) { // From the first proton
	  Parton parton(mScatter.at(0)); // Copy pdfs, proton spin, x
	  // momentum and species can change after scattering, so update
	  parton.SetPxPyPz(TVector3(px, py, pz));
	  parton.type = type;
	  partons.push_back(parton);
	} // if
	else if(7 == parent) { // From the second proton
	  Parton parton(mScatter.at(1));
	  // momentum and species can change after scattering, so update
	  parton.SetPxPyPz(TVector3(px, py, pz));
	  parton.type = type;
	  partons.push_back(parton);
	} // else if
      } // if
         
      // Also accumulate J/Psi in the hadron list
      // \todo Add config flag option
      if((443 == type) and (1 == status)) {
	Hadron hadron;
	hadron.px =  px; 
	hadron.py =  py; 
	hadron.pz =  pz; 
	hadron.z = 1; 
	hadron.lasthad = 0; 
	hadron.type = JPSI; 
	hadron.mParton = Parton(0., 0., 0., 21, -1, 0, 0); // gluon 
	AddHadron(hadron);
      } // if
      if (mtparticleout > 0 ){
	//add to the tparticles array
	Int_t daug1 = particle->GetFirstChild() - 1;
	Int_t daug2 = particle->GetLastChild() - 1;
	Int_t mother1 = particle->GetParent() - 1;
	Int_t mother2 = -1;
	Double_t E = particle->GetEnergy();

	Double_t vx = particle->GetVx();
	Double_t vy = particle->GetVy();
	Double_t vz = particle->GetVz();
	Double_t t = particle->GetTime();
	   
	TLorentzVector p(px, py, pz, E);
	TLorentzVector v(vx, vy, vz, t);
	TParticle tparticle(type, status, mother1, mother2, daug1, daug2, p, v);
	//add the spin states to the protons and scattering partons
	//so far only polarized in the y direction
	if ( counter == 0 || counter == 1){ 
	  tparticle.SetPolarisation(0,mScatter.at(counter).pspin,0);
	}
	if ( counter == 4 || counter == 5){ 
	  tparticle.SetPolarisation(0,mScatter.at(counter-4).spin,0);
	}
	AddParticle(tparticle);
      }
      ++counter;
    } // while
    partons.sort(SortPartonsByGroup());
      BOOST_FOREACH(const Parton& p, partons) {
	AddParton(p);
      } // BOOST_FOREACH
  }

  double Event::ComputeTransverseSpin(const Parton& parton, Distributions& pdfs) {
    RandomSpin random(pdfs.Unpolarised().get(),
		      pdfs.Transversity().get());
    return random.Generate(parton, mQSquared);
  }

  // Compute the ratio of the Sivers distribution to the unpolarised.
  double Event::ComputeSiversWeight(const Parton& parton, Distributions& pdfs) {
    PartonContent unpol;
    PartonContent sivers;
    pdfs.Unpolarised()->Evaluate(parton, mQSquared, unpol);
    pdfs.Sivers()->Evaluate(parton, mQSquared, sivers);
    int type = parton.Lhapdf();
    // Sivers is bounded by unpolarised
    double weight = sivers[type] / unpol[type];
    if(isnan(weight)) {
      weight = 0.;
    } // if
    return weight;
  }

  double Event::ComputeSiversSpinFactor(const Parton& parton) {
    if(parton.Pt() > 0.) {
      // SPECIAL NOTE: The sign of the Sivers spin factor has a minus sign relative to 
      // the extractions done by the Torino group because they used the Trento
      // convention to extract the Sivers distributions, where Sivers>0 corresponds
      // to the preference for a quark to move to the observers left if the proton is 
      // moving towards the observer and the spin is pointing up. pp measurements
      // use the opposite convention, and hence the minus sign. 
      /** \todo Reverse this, as I redefined the sivers such that it
	  follows the pp convention i.e. sivers > 0 gives preference to
	  observer's right*/
      return parton.pspin * (parton.px / parton.Pt());
    } // if
    else {
      return 0.;
    } // else
  }

  std::vector<tppmc::Parton*> Event::Partons() const {
    std::vector<Parton*> partons;
    TIter next(&mPartons);
    Parton* parton(NULL);
    while((parton = dynamic_cast<Parton*>(next()))) {
      partons.push_back(parton);
    } // while
    return partons;
  }

  std::vector<tppmc::Hadron*> Event::Hadrons() const {
    std::vector<Hadron*> hadrons;
    TIter next(&mHadrons);
    Hadron* hadron(NULL);
    while((hadron = dynamic_cast<Hadron*>(next()))) {
      hadrons.push_back(hadron);
    } // while
    return hadrons;
  }

  std::vector<TParticle*> Event::Particles() const {
    std::vector<TParticle*> Particles;
    TIter next(&mParticle);
    TParticle* tparticle(NULL);
    while((tparticle = dynamic_cast<TParticle*>(next()))) {
      Particles.push_back(tparticle);
    } // while
    return Particles;
  }

  unsigned Event::AddParton(const Parton& parton) {
    const unsigned size = mPartons.GetEntries();
    new(mPartons[size]) Parton(parton);
    return mPartons.GetEntries();
  }

  unsigned Event::AddHadron(const Hadron& hadron) {
    const unsigned size = mHadrons.GetEntries();
    new(mHadrons[size]) Hadron(hadron);
    return mHadrons.GetEntries();
  }

  unsigned Event::AddParticle(const TParticle& tparticle) {
    const unsigned size = mParticle.GetEntries();
    new(mParticle[size]) TParticle(tparticle);
    return mParticle.GetEntries();
  }

  int Event::Fragment(FragmentationFunction& ff) {
    unsigned initialSize = mHadrons.GetEntries();
    BOOST_FOREACH(Parton* p, Partons()) {
      FragmentParton(*p, ff);
    } // BOOST_FOREACH
    return mHadrons.GetEntries() - initialSize;
  }

  int Event::FragmentParton(Parton& parton,
			    FragmentationFunction& ff) {
    unsigned initialSize = mHadrons.GetEntries();
    const double pTotalIntial = parton.PxPyPz().Mag();
    // We update the z range after each hadron is generated
    // so that we don't exceed the remaining momentum that
    // the parton has left.
    double zMax = ff.ZMax();
    // Also track the remaining momentum to calculate z.
    double pTotal = pTotalIntial;
      
    // Generate the rotation matrix that will put us into the 
    // fragmenting parton frame:
    TRotation rotM; 
    rotM.SetZAxis(parton.PxPyPz());
    TRotation rotMInv = rotM.Inverse();
      
    // We generate fragments until we're below a limit, so
    // keep a running sum of the total hadron momentum generated.
    Hadron runsum;
    runsum.px = 0.0;
    runsum.py = 0.0;
    runsum.pz = 0.0; 
      
    // Track the location of the spin axis: 
    const TVector3 spin_lab(0., double(parton.Spin()), 0.); // Y direction
    TVector3 spin_parton_frame = rotMInv * spin_lab; 
      
    // What is the angle between the spin axis in the parton frame
    // and the plane parallel to the parton momentum?
    // i.e. effectively no polarisation when the parton is parallel
    // to the spin direction.
    double cos_spin_depol = spin_parton_frame.Perp();
    parton.cos_spin_depol = cos_spin_depol;
      
    // Get the azimuthal angle of the spin axis in the new frame
    double phi_spin = TMath::ATan2(spin_parton_frame.Y(),
				   spin_parton_frame.X());
    parton.phi_spin = phi_spin;
     if(parton.PxPyPz().Mag() > sqrt(ff.MinimumScale())) {
      do {
	// Use the parton pT^2 as the fragmentation scale.
	// If the parton pT is below the lower limit, use that.
	const double pt2 = std::max(pow(parton.Pt(), 2.),
				    ff.MinimumScale());
	ff.SetZRange(ff.ZMin(), zMax);
	std::vector<Hadron> hadrons = ff.Fragment(parton, pt2);
	// Stop once we don't generate any more hadrons.
	if(hadrons.empty()) {
	  break;
	} // if
	BOOST_FOREACH(Hadron& h, hadrons) {
	  // Hadrons are in the parton frame.
	  // Do some further calculations and rotate to lab frame.
	  double newPz = h.Z() * pTotalIntial;
	  pTotal -= newPz;
	  double jt = h.PxPyPz().Pt();
	  // Careful!
	  // Phi we got from fragmentation is phi wrt spin direction.
	  // The sense of the angle phi is opposite the sense
	  // in the xy plane, hence the minus sign.
	  // Also, the phi angle from fragmentation is the angle between
	  // spin and (quark x hadron) momenta, so the hadron momentum is
	  // at 90 degrees to this, hence additional pi/2 angle change.
	  double phi = h.PxPyPz().Phi();
	  h.collinsangle = phi;
	  double new_px = jt * cos(phi_spin - TMath::PiOver2() - phi);
	  double new_py = jt * sin(phi_spin - TMath::PiOver2() - phi);
	  runsum.px += new_px;
	  runsum.py += new_py;
	  runsum.pz += newPz;
               
	  TVector3 fragment(new_px, new_py, newPz);
	  fragment.Transform(rotM); 
               
	  h.px = fragment.X();
	  h.py = fragment.Y();
	  h.pz = fragment.Z();
	  h.lasthad = 0; 

	  Int_t daug1 =  - 1;
	  Int_t daug2 =  - 1;
	  Int_t mother1 = h.Group()+5; // to adjust it to the outgoing partons
	  // need to check, that order did not change wrt to initial particles
	  Int_t mother2 = -1;


	  Double_t px = h.px;
	  Double_t py = h.py;
	  Double_t pz = h.pz;
	  //	 Double_t E = h.GetEnergy(); //not available
	  double mass = (abs(h.type)==211)?0.13957018:abs(h.type)==111?0.1349766:abs(h.type)==321?0.493677:abs(h.type)==221?0.547852:abs(h.type)==13?0.105658367:0;
			
	  Double_t E = sqrt (mass*mass + px*px + py*py + pz*pz );

	  if (h.z>1.0)
	    printf("warning hadron has z larger than one %.3f, type %d, parton energy %.2f\n",h.z,h.type,E/parton.P());
	   

	  TLorentzVector p(px, py, pz, E);
	  Double_t vx = 0;
	  Double_t vy = 0;
	  Double_t vz = 0;
	  Double_t t = 0;
	  TLorentzVector v(vx, vy, vz, t);
	  TParticle tparticle(h.type, h.lasthad, mother1, mother2, daug1, daug2, p, v);

	  if(h.PxPyPz().Eta() > mEtaMin and h.PxPyPz().Eta() < mEtaMax and
	     h.PxPyPz().Pt() > mMinPt) {
	    if (mtparticleout >0){
	      AddParticle(tparticle);
	    }
	    if (mtparticleout %2 ==0 ){
	      AddHadron(h);
	    }
	  } // if
	  zMax = pTotal / pTotalIntial; 
	} // BOOST_FOREACH
      } while(pTotal > 0.5 and zMax > ff.ZMin());
    } // if
    assert((pTotalIntial - pTotal - runsum.pz)<1.e-4);

    // Conserver any remaining parton momentum with a single soft hadron
    // Skip this if we aren't retaining the event (nhadrons == 0) or there
    // is an explicity eta cut (eta > -999).
    /**
       \todo We could potentially help charge non-conservation a little here by
       forcing the last hadron to have opposite charge to the sum of the other
       hadrons e.g if we have an excess of pi+ from the other fragmentations,
       force pi- and vice versa.
       \todo need to check momentum conservation works properly
    */
    if(pTotal > 0. and Hadrons().size() > 0 and mEtaMin<=-999. and mEtaMax>=999.) {
      // Running sum is in parton frame, transform difference 
      // back to lab frame - conserves momentum of fragmenting parton 
      TVector3 fragment(-1.0*runsum.px, -1.0*runsum.py, pTotal); 
      fragment.Transform(rotM); 
      Hadron hadron;
      hadron.px =  fragment.X();
      hadron.py =  fragment.Y();
      hadron.pz =  fragment.Z();
      hadron.z = pTotal / pTotalIntial;
      hadron.lasthad = 1;
      hadron.type = 111; // pi0
      hadron.mParton = parton;
      hadron.phigen = gRandom->Uniform(-TMath::Pi(), TMath::Pi());
      //calculate angle between hadron momentum and y-component of parton spin, projected into the plane perpendicular to the parton momentum
      TVector3 parton_mom = parton.PxPyPz();
      parton_mom = parton_mom.Unit();
      TVector3 spin_vec(0., parton.spin, 0.);
      TVector3 h_mom(hadron.px, hadron.py, hadron.pz);
      double projection_val = h_mom.Dot(parton_mom);
      TVector3 h_projection = h_mom - (projection_val * parton_mom);
      projection_val = spin_vec.Dot(parton_mom);
      TVector3 spin_projection = spin_vec - (projection_val * parton_mom);
      double cos_angle = (spin_projection.Unit()).Dot( h_projection.Unit() );
      hadron.collinsangle = acos(cos_angle);
      if (mtparticleout %2 >0 ){
	AddHadron(hadron);
      }
	 
      Int_t daug1 =  - 1;
      Int_t daug2 =  - 1;
      Int_t mother1 = hadron.Group()+5; // to adjust it to the outgoing partons
      // need to check, that order did not change wrt to initial particles
      Int_t mother2 = -1;


      Double_t px = hadron.px;
      Double_t py = hadron.py;
      Double_t pz = hadron.pz;
      //	 Double_t E = hadron.GetEnergy(); //not available
      double mass = (abs(hadron.type)==211)?0.13957018:abs(hadron.type)==111?0.1349766:abs(hadron.type)==321?0.493677:abs(hadron.type)==221?0.547852:abs(hadron.type)==13?0.105658367:0;
			
      Double_t E = sqrt (mass*mass + px*px + py*py + pz*pz );

      if (hadron.z>1.0)
	printf("warning last  hadron has z larger than one %.3f, type %d, parton energy %.2f\n",hadron.z,hadron.type,E/parton.P());
	   
      if (mtparticleout>0){
	TLorentzVector p(px, py, pz, E);
	Double_t vx = 0;
	Double_t vy = 0;
	Double_t vz = 0;
	Double_t t = 0;
	TLorentzVector v(vx, vy, vz, t);
	TParticle tparticle(hadron.type, hadron.lasthad  , mother1, mother2, daug1, daug2, p, v);
	AddParticle(tparticle);
      }
    } // if
    return mHadrons.GetEntries() - initialSize;
  }

  Hadron Event::HadronSum(int group) const {
    Hadron sum;
    BOOST_FOREACH(Hadron* hadron, Hadrons()) {
      if(hadron->Group() == group or group == 0) {
	sum += *hadron;
      } // if
    } // BOOST_FOREACH
    return sum;
  }

  unsigned Event::Multiplicity(unsigned n) const {
    unsigned mult(0);
    BOOST_FOREACH(Hadron* hadron, Hadrons()) {
      if(unsigned(hadron->Group()) == n) ++mult;
    } // BOOST_FOREACH
    return mult;
  }

  //
  // class EventFactory
  //
   
  Event* EventFactory::Create(double rotation) {
      
    Event* event = new Event(rotation, -999., 999.);
      
    return event;
  }
} // namespace tppmc
